// winget.cpp : Defines the entry point for the console application.
//

#include "stdafx.h"
#include <stdio.h>
#include <windows.h>
#include <bits.h>
#include <strsafe.h>

IBackgroundCopyManager *g_pbcm = NULL;

class CNotifyInterface : public IBackgroundCopyCallback {
    LONG m_lRefCount;
    HANDLE m_hEvent;

public:
    CNotifyInterface();
    CNotifyInterface(HANDLE hEvent);
    ~CNotifyInterface();

    // IUnknown
    HRESULT __stdcall QueryInterface(REFIID riid, LPVOID *ppvObj);
    ULONG __stdcall AddRef();
    ULONG __stdcall Release();

    // IBackgroundCopyCallback methods
    HRESULT __stdcall JobTransferred(IBackgroundCopyJob *pJob);
    HRESULT __stdcall JobError(IBackgroundCopyJob *pJob, IBackgroundCopyError *pError);
    HRESULT  __stdcall JobModification(IBackgroundCopyJob *pJob, DWORD dwReserved);
};

CNotifyInterface::CNotifyInterface()
{
    m_lRefCount = 1;
    m_hEvent = INVALID_HANDLE_VALUE;
}

CNotifyInterface::CNotifyInterface(
    HANDLE hEvent
    )
{
    m_lRefCount = 1;
    m_hEvent = hEvent;
}

CNotifyInterface::~CNotifyInterface()
{
}

HRESULT
CNotifyInterface::QueryInterface(
    REFIID riid,
    LPVOID *ppvObj
    )
{
  if (riid == __uuidof(IUnknown) || riid == __uuidof(IBackgroundCopyCallback)) {
    *ppvObj = this;
  } else {
    *ppvObj = NULL;
    return E_NOINTERFACE;
  }

  AddRef();
  return NOERROR;
}

ULONG 
CNotifyInterface::AddRef() 
{
    return InterlockedIncrement(&m_lRefCount);
}

ULONG 
CNotifyInterface::Release() 
{
    ULONG  ulCount = InterlockedDecrement(&m_lRefCount);

    if(0 == ulCount) {
        delete this;
    }

    return ulCount;
}

HRESULT 
CNotifyInterface::JobTransferred(
    IBackgroundCopyJob *pJob
    )
{
    HRESULT hr;

    hr = pJob->Complete();
    if (FAILED(hr)) {
        // TODO: Handle error here
    }
    
    if (m_hEvent != INVALID_HANDLE_VALUE) {
        SetEvent(m_hEvent);
    }

    return S_OK;
}

HRESULT 
CNotifyInterface::JobError(
    IBackgroundCopyJob *pJob, 
    IBackgroundCopyError *pError
    )
{
    HRESULT hr;
    BG_ERROR_CONTEXT context;
    HRESULT errorCode = S_OK;

    hr = pError->GetError(&context, &errorCode);

    // TODO: Handle error here

    if (m_hEvent != INVALID_HANDLE_VALUE) {
        SetEvent(m_hEvent);
    }

    return S_OK;
}

HRESULT
CNotifyInterface::JobModification(
    IBackgroundCopyJob *pJob, 
    DWORD dwReserved
    )
{
    HRESULT hr;
    BG_JOB_PROGRESS progress;
    BG_JOB_STATE state;

    hr = pJob->GetProgress(&progress);

    if (SUCCEEDED(hr)) {
        hr = pJob->GetState(&state);
    }

    if (SUCCEEDED(hr)) {
        // TODO: Print progress report for user
    }

    return S_OK;
}

HRESULT
CreateDownloadJob(
    __in PWSTR Url,
    __in PWSTR Dest,
    __in BOOL Async
	)
{
    HRESULT hr = S_OK;
    GUID jobId = {0};
    IBackgroundCopyJob *pJob = NULL;
    CNotifyInterface *pNotify = NULL;
    HANDLE hEvent = INVALID_HANDLE_VALUE;

    hr = CoCreateInstance(__uuidof(BackgroundCopyManager),
                          NULL,
                          CLSCTX_LOCAL_SERVER,
                          __uuidof(IBackgroundCopyManager),
                          (void**)&g_pbcm);

    if (SUCCEEDED(hr)) {
        hr = g_pbcm->CreateJob(L"WingetDownloadJob", BG_JOB_TYPE_DOWNLOAD, &jobId, &pJob);
    }

    if (SUCCEEDED(hr)) {
        hr = pJob->SetPriority(BG_JOB_PRIORITY_FOREGROUND);
    }

    if (!Async) {
        hEvent = CreateEvent(NULL, TRUE, FALSE, NULL);
        pNotify = new CNotifyInterface(hEvent);
    } else {
        pNotify = new CNotifyInterface();
    }

    if (SUCCEEDED(hr)) {
        hr = pJob->SetNotifyInterface(pNotify);
    }

    if (SUCCEEDED(hr)) {
        hr = pJob->SetNotifyFlags(BG_NOTIFY_JOB_TRANSFERRED | BG_NOTIFY_JOB_ERROR | BG_NOTIFY_JOB_MODIFICATION);
    }

    if (SUCCEEDED(hr)) {
        hr = pJob->AddFile(Url, Dest);
    }

    if (SUCCEEDED(hr)) {
        hr = pJob->Resume();
    }

    if (!Async) {
        // Wait for the download to complete
        WaitForSingleObject(hEvent, INFINITE);
    }

    if (hEvent != INVALID_HANDLE_VALUE) {
        CloseHandle(hEvent);
    }

    return hr;
}

void
PrintUsage()
{
    printf("%s\n", "winget.exe: Usage information");
    printf("%s\n", "winget.exe \"<URL>\"");
}

int
_tmain(
    int argc, 
    _TCHAR* argv[])
{
	HRESULT hr = S_OK;
	BOOL comInitialized = FALSE;
    WCHAR destPath[MAX_PATH];
    WCHAR *p;

    if (argc != 2) {
        PrintUsage();
        return 1;
    }

    if (lstrlen(argv[1]) > 1000) {
        printf("Invalid URL.\n");
        return 1;
    }

    p = wcsrchr(argv[1], L'/');
    if (p == NULL) {
        printf("Invalud URL.\n");
        return 1;
    }

    if (GetCurrentDirectory(MAX_PATH, destPath) == 0) {
        printf("Could not retrieve current directory path.\n");
        return GetLastError();
    }
	
    if (FAILED(StringCchCat(destPath, MAX_PATH, L"\\")) ||
        FAILED(StringCchCat(destPath, MAX_PATH, p))) {
            printf("Could not determine download destination path.\n");
            return 1;
    }

    printf("Download URL: %ws\n", argv[1]);
    printf("Destination path: %ws\n", destPath);

	//
	// BITS is accessed using a COM object, so initialize COM for the duration
	// of this process.
	//
	hr = CoInitializeEx(NULL, COINIT_MULTITHREADED);
	if (SUCCEEDED(hr)) {
		comInitialized = TRUE;
	}

    hr = CreateDownloadJob(argv[1], destPath, FALSE);
    if (FAILED(hr)) {
        printf("Download job failed with error 0x%08x\n", hr);
        return 1;
    }

	if (comInitialized) {
		CoUninitialize();
	}
	return 0;
}

